/*
 * Copyright 2011 Noctarius (Christoph Engelbert)
 * Part of www.sourceprojects.org 
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may 
 * not use this file except in compliance with the License. 
 * You may obtain a copy of the License at 
 *    http://www.apache.org/licenses/LICENSE-2.0 
 *
 * Unless required by applicable law or agreed to in writing, software 
 * distributed under the License is distributed on an "AS IS" BASIS, 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 * See the License for the specific language governing permissions and 
 * limitations under the License.
 */
package org.sourceprojects.cappadocia;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.StringTokenizer;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.jar.Attributes;
import java.util.jar.JarFile;
import java.util.jar.Manifest;
import java.util.zip.ZipEntry;
import java.util.zip.ZipFile;

import org.apache.log4j.Logger;
import org.objectweb.asm.ClassReader;

public class ClassPool {
	private static final Logger LOG = Logger.getLogger(ClassPool.class);

	private final Map<String, File> locations = new LinkedHashMap<String, File>();
	private final Lock locationsLock = new ReentrantLock();

	private final Map<String, ClassInfo> classCache = new HashMap<String, ClassInfo>();
	private final Lock classCacheLock = new ReentrantLock();

	private final AtomicLong lastUpdate = new AtomicLong();
	private final AtomicLong lastInsert = new AtomicLong();

	private final AtomicBoolean dynamicResolving = new AtomicBoolean(true);

	public void addClasspath() throws IOException {
		addClasspath(false);
	}

	public void addClasspath(final boolean includeRuntimeEnvironment)
			throws IOException {

		final String classpath = System.getProperty("java.class.path");
		final StringTokenizer tokenizer = new StringTokenizer(classpath,
				File.pathSeparator);

		while (tokenizer.hasMoreTokens()) {
			final String segment = tokenizer.nextToken();

			if (!isRuntimeJar(segment)) {
				addFile(new File(segment));
			}
		}

		if (includeRuntimeEnvironment) {
			addLibraryClasspath();
		}
	}

	public boolean addFile(final File path) throws IOException {
		final String absolutePath = path.getAbsolutePath();

		locationsLock.lock();
		try {
			if (!locations.containsKey(absolutePath)) {
				locations.put(absolutePath, path);
				lastInsert.set(System.currentTimeMillis());

				if (isClassPool(path)) {
					if (isManifestContaining(path)) {
						findManifestLocations(path);
					}
				}
			}
		} finally {
			locationsLock.unlock();
		}

		return false;
	}

	public void refreshClassCache() throws IOException {
		classCacheLock.lock();
		try {
			if (needsRefresh()) {
				classCache.clear();
				final Collection<ClassInfo> classes = processLocations();
				for (final ClassInfo info : classes) {
					classCache.put(info.getCanonicalName(), info);
				}

				lastUpdate.set(System.currentTimeMillis());
			}

		} finally {
			classCacheLock.unlock();
		}
	}

	public boolean isDynamicResolving() {
		return dynamicResolving.get();
	}

	public void setDynamicResolving(final boolean dynamicResolving) {
		this.dynamicResolving.set(dynamicResolving);
	}

	public boolean needsRefresh() {
		return lastInsert.longValue() > lastUpdate.longValue();
	}

	public Collection<ClassInfo> getClassCache() {
		return Collections.unmodifiableCollection(classCache.values());
	}

	public ClassInfo getClassInfo(final String canonicalName) {
		classCacheLock.lock();
		try {
			try {
				refreshClassCache();
			} catch (final IOException e) {
				LOG.warn("ClassCache-Refresh failed", e);
			}

			final ClassInfo classInfo = classCache.get(canonicalName);

			if (classInfo != null) {
				return classInfo;
			}

			if (dynamicResolving.get()) {
				try {
					final ClassReader classReader = new ClassReader(
							canonicalName);
					final ClassInfoImpl result = new ClassInfoImpl(this);
					classReader.accept(new ClassInfoParser(this, result), 0);

					classCache.put(canonicalName, result);

					return result;

				} catch (final IOException e) {
					LOG.warn("Reading class from classpath failed", e);
				}
			}

			throw new ClassQueryException("Interface '" + canonicalName
					+ "' could not be resolved. Please "
					+ "ensure that it is in search classpool.");

		} finally {
			classCacheLock.unlock();
		}
	}

	@Override
	public int hashCode() {
		final int prime = 31;
		int result = 1;
		result = prime * result
				+ ((locations == null) ? 0 : locations.hashCode());
		return result;
	}

	@Override
	public boolean equals(final Object obj) {
		if (this == obj) {
			return true;
		}
		if (obj == null) {
			return false;
		}
		if (getClass() != obj.getClass()) {
			return false;
		}
		final ClassPool other = (ClassPool) obj;
		if (locations == null) {
			if (other.locations != null) {
				return false;
			}
		} else if (!locations.equals(other.locations)) {
			return false;
		}
		return true;
	}

	@Override
	public String toString() {
		final StringBuilder builder = new StringBuilder();
		builder.append("ClassPool [locations=").append(locations).append("]");
		return builder.toString();
	}

	private boolean isManifestContaining(final File path) {
		final String filename = path.getName().toLowerCase();

		return filename.endsWith(".jar") || filename.endsWith(".war");
	}

	private void findManifestLocations(final File file) throws IOException {
		final JarFile jar = new JarFile(file);
		final Manifest manifest = jar.getManifest();

		if (manifest == null) {
			return;
		}

		final Attributes attributes = manifest.getMainAttributes();
		for (final Object key : attributes.keySet()) {
			final String attribute = key.toString();
			final String value = (String) attributes.get(key);

			if ("Class-Path".equals(attribute)) {
				final StringBuilder builder = new StringBuilder();
				final StringTokenizer tokenizer = new StringTokenizer(value);
				while (tokenizer.hasMoreTokens()) {
					builder.setLength(0);
					final String element = tokenizer.nextToken();
					final String parent = file.getParent();

					if (parent != null) {
						builder.append(parent);
						builder.append(File.pathSeparator);
					}

					builder.append(element);
				}

				addFile(new File(builder.toString()));
			}
		}
	}

	private void addLibraryClasspath() throws IOException {
		final String libPath = findLibPath();
		final StringTokenizer tokenizer = new StringTokenizer(libPath,
				File.pathSeparator);

		while (tokenizer.hasMoreTokens()) {
			final String segment = tokenizer.nextToken();
			addFile(new File(segment));
		}
	}

	private String findLibPath() throws IOException {
		final StringBuilder libPath = new StringBuilder();
		final String basePath = new StringBuilder(
				System.getProperty("java.home")).append(File.separator)
				.append("lib").toString();

		final File path = new File(basePath);
		final List<File> files = recursiveAddFiles(path);
		for (final File file : files) {
			if (file.exists()) {
				libPath.append(File.pathSeparator);
				libPath.append(file.getAbsolutePath());
			}
		}

		final String sunBootLibPath = System.getProperty("sun.boot.class.path");
		if (sunBootLibPath != null) {
			libPath.append(File.pathSeparator);
			libPath.append(sunBootLibPath);
		}

		return libPath.toString();
	}

	private List<File> recursiveAddFiles(final File basePath)
			throws IOException {

		final List<File> files = new ArrayList<File>();
		if (basePath.exists() && basePath.isDirectory()) {
			for (final File child : basePath.listFiles()) {
				if (child.isDirectory()) {
					files.addAll(recursiveAddFiles(child));

				} else {
					if (isClassPool(child)) {
						files.add(child);
					}
				}
			}
		}

		return files;
	}

	private boolean isRuntimeJar(final String path) {
		for (final String entry : ClassPoolConstants.RUNTIME_ENTRIES) {
			if (path.endsWith(entry)) {
				return true;
			}
		}

		return false;
	}

	private boolean isClassPool(final File path) {
		final String filename = path.getName().toLowerCase();

		if (path.exists()) {
			return path.isDirectory() || filename.endsWith(".zip")
					|| filename.endsWith(".jar") || filename.endsWith(".war")
					|| filename.endsWith(".ear");
		}

		return false;
	}

	private Set<ClassInfo> processLocations() throws IOException {
		locationsLock.lock();
		try {
			final Set<ClassInfo> classInfos = new HashSet<ClassInfo>();
			for (final File entry : locations.values()) {
				if (entry.exists()) {
					if (entry.isDirectory()) {
						classInfos.addAll(processDirectory(entry));

					} else if (entry.getName().toLowerCase().endsWith(".jar")) {
						classInfos.addAll(processJarFile(entry));

					} else if (entry.getName().toLowerCase().endsWith(".zip")) {
						classInfos.addAll(processZipFile(entry));

					}
				}
			}

			return classInfos;

		} finally {
			locationsLock.unlock();
		}
	}

	private Set<ClassInfo> processJarFile(final File entry) throws IOException {
		return processCompressedFolder(new JarFile(entry), entry);
	}

	private Set<ClassInfo> processZipFile(final File entry) throws IOException {
		return processCompressedFolder(new ZipFile(entry), entry);
	}

	private Set<ClassInfo> processDirectory(final File path) throws IOException {
		final Set<ClassInfo> classInfos = new HashSet<ClassInfo>();

		for (final File child : path.listFiles()) {
			if (child.isDirectory()) {
				classInfos.addAll(processDirectory(child));

			} else if (child.getName().toLowerCase().endsWith(".class")) {
				classInfos.add(visitClass(new FileInputStream(child)));
			}
		}

		return classInfos;
	}

	private Set<ClassInfo> processCompressedFolder(final ZipFile zipFile,
			final File entry) throws IOException {

		final Set<ClassInfo> classInfos = new HashSet<ClassInfo>();
		final Enumeration<? extends ZipEntry> enumeration = zipFile.entries();
		while (enumeration.hasMoreElements()) {
			final ZipEntry zipEntry = enumeration.nextElement();

			if (!zipEntry.isDirectory()
					&& zipEntry.getName().toLowerCase().endsWith(".class")) {

				classInfos.add(visitClass(zipFile.getInputStream(zipEntry)));
			}

		}

		return classInfos;
	}

	private ClassInfo visitClass(final InputStream inputStream)
			throws IOException {

		try {
			final ClassInfoImpl classInfo = new ClassInfoImpl(this);
			final ClassReader classReader = new ClassReader(inputStream);
			classReader.accept(new ClassInfoParser(this, classInfo), 0);
			return classInfo;

		} finally {
			inputStream.close();
		}
	}

	ClassInfo resolveClass(final String canonicalName) {
		classCacheLock.lock();
		try {
			try {
				refreshClassCache();
			} catch (final IOException e) {
				LOG.warn("ClassCache-Refresh failed", e);
			}

			final ClassInfo classInfo = classCache.get(canonicalName);

			if (classInfo != null) {
				return classInfo;
			}

			if (dynamicResolving.get()) {
				try {
					final ClassReader classReader = new ClassReader(
							canonicalName);
					final ClassInfoImpl result = new ClassInfoImpl(this);
					classReader.accept(new ClassInfoParser(this, result), 0);

					classCache.put(canonicalName, result);

					return result;

				} catch (final IOException e) {
					LOG.debug("IOException while dynamic resolving", e);
				}
			}

			LOG.debug("Dynamic resolving of class '" + canonicalName
					+ "'from classpath failed");
			return null;

		} finally {
			classCacheLock.unlock();
		}
	}

}
